cuda-gen: Add CUDA Graph capture and replay for composite operators - #1987
cuda-gen: Add CUDA Graph capture and replay for composite operators#1987Nafees01 wants to merge 18 commits into
Conversation
Are there specific operators you were unable to use in Ratel? None of the QFunctions there should alloc during operator application. |
|
There's some style errors here, be sure to |
zatkins-dev
left a comment
There was a problem hiding this comment.
I'm sure you're still working on this, but here's an initial review to correct some high-level design issues and removed functionality that I think should be fixed.
|
I would be interested to see if this speeds up e.g. /gpu/cuda/ref, when recording the graph over the entire |
|
Note - merging main into your branch instead of rebasing means I'd recommend we squash-merge this PR |
The only example where the contact (Nitsche) sub-operator couldn't be captured was the quasistatic contact case (ex02-quasistatic-elasticity-linear-platen). Everything else I tried static, dynamic, multi-material, neo-Hookean, Mooney-Rivlin captured and replayed fine. And you're right that the QFunction doesn't allocate; the cudaMalloc comes from libCEED's own device-side handling of the contact data, not the QFunction itself. It runs either way correctly; the contact operators just fall back to a normal apply rather than being replayed. |
This is a pretty big red flag to me. The Nitsche contact operators don't actually do anything out of the ordinary on the libCEED side. Them failing to work means that something about this setup (very possibly something with passive inputs and outputs) is violating the libCEED interface assumptions. |
Thanks for flagging this. On re-checking linear-platen with CEED_DEBUG=1, the contact operators appear to fail earlier than graph capture: cuda-gen hits an NVRTC compile error (GradTransposeTensor3dFlattened defined twice in cuda-shared-basis-tensor-flattened-templates.h) and then falls back to /gpu/cuda/ref. So this looks like a pre-existing cuda-gen JIT issue triggered by the contact operator's basis configuration, rather than something specific to Nitsche or CUDA graphs. |
Sweet, good catch. I'll fix that separately. |
| @brief Set CUDA function pointer to evaluate action at quadrature points | ||
|
|
||
| If the backend does not support `CUfunction` pointers for QFunctions, the call succeeds without effect. | ||
| When unsupported, a message is emitted via `CeedDebug`. |
There was a problem hiding this comment.
| @brief Set CUDA function pointer to evaluate action at quadrature points | |
| If the backend does not support `CUfunction` pointers for QFunctions, the call succeeds without effect. | |
| When unsupported, a message is emitted via `CeedDebug`. | |
| @brief Set CUDA function pointer to evaluate action at quadrature points. | |
| If the backend does not support `CUfunction` pointers for QFunctions, then the call succeeds without effect. | |
| When unsupported, a message is emitted via `CeedDebug`. |
minor
| @brief Enable or disable CUDA Graph capture/replay for a `CeedOperator` | ||
|
|
||
| If the backend does not support CUDA Graphs for operators, the call succeeds without effect. | ||
| When unsupported, a message is emitted via `CeedDebug`. |
There was a problem hiding this comment.
| @brief Enable or disable CUDA Graph capture/replay for a `CeedOperator` | |
| If the backend does not support CUDA Graphs for operators, the call succeeds without effect. | |
| When unsupported, a message is emitted via `CeedDebug`. | |
| @brief Enable or disable CUDA Graph capture/replay for a `CeedOperator`. | |
| If the backend does not support CUDA Graphs for operators, then the call succeeds without effect. | |
| When unsupported, a message is emitted via `CeedDebug`. |
Summary
Adds CUDA Graph capture and replay to the 'cuda-gen' backend for composite operators. On the second apply, the kernel sequence is captured into a cudaGraph; all subsequent applies replay the instantiated graph. QFunction contexts are re-synced to device before each replay so time- and load-dependent parameters stay correct.
Sub-operators that cannot be captured, for example, contact operators that allocate memory (cudaMalloc) during apply, which CUDA doesn't allow inside a capture. When this happens, the capture is properly closed to leave the stream in a clean state, any leftover CUDA error is cleared, and that operator falls back to running normally on every call.
Validated against both /gpu/cuda/ref and cuda-gen without graphs across static, quasistatic, contact, multi-material, and dynamic elasticity examples. Strain energy and displacements match to ~1e-13 in every case, with zero graph fallbacks on capturable operators.
Performance-wise, graphs are within ±1% of cuda-gen without graphs across all sizes, so no regression, but no real speedup either. The reason is that cuda-gen already fuses everything into a few large kernels, so there's very little launch overhead left for graphs to remove.
Validated correctness against /gpu/cuda/ref and cuda-gen (CEED_DISABLE_GRAPH=1) on the following Ratel examples:
Strain energy matches to ~1e-13 in all cases. Zero graph fallbacks on capturable operators.